#!/usr/bin/env python3
"""
FASTA文件合并工具 - 将多行序列合并为单行格式

功能：
- 将多行序列合并为单行格式
- 支持批量处理多个文件
- 提供文件名前缀添加功能
- 支持FASTA格式验证

作者: 陈振玺
版本: 0.5.0.20240803
创建时间: 2024年8月
"""

import argparse
import os
import sys
from pathlib import Path
from typing import List, Dict, Optional, Tuple

from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord


class FastaMergerError(Exception):
    """FASTA合并器基础异常类"""
    pass


class FileNotFoundError(FastaMergerError):
    """文件未找到异常"""
    pass


class InvalidFastaFormatError(FastaMergerError):
    """无效FASTA格式异常"""
    pass


class FastaMerger:
    """
    FASTA文件合并器 - 将多行序列合并为单行格式
    
    该类提供专业的FASTA文件处理功能，基于BioPython库开发，
    用于将多行序列合并为单行格式，支持批量处理和格式验证。
    
    属性:
        processed_count (int): 成功处理的序列数量
        error_count (int): 处理失败的序列数量
    """
    
    def __init__(self):
        """初始化FASTA合并器"""
        self.processed_count = 0
        self.error_count = 0
        
    def merge_sequences(self, input_file: str, output_file: str, 
                     add_filename_prefix: bool = False) -> bool:
        """
        合并单个FASTA文件中的多行序列为单行格式
        
        参数:
            input_file (str): 输入FASTA文件路径
            output_file (str): 输出FASTA文件路径
            add_filename_prefix (bool): 是否添加文件名前缀到序列ID
            
        返回:
            bool: 处理成功返回True，失败返回False
            
        异常:
            FileNotFoundError: 输入文件不存在
            InvalidFastaFormatError: FASTA格式无效
            PermissionError: 输出文件无写入权限
        """
        try:
            # 检查输入文件是否存在
            if not os.path.exists(input_file):
                raise FileNotFoundError(f"输入文件不存在: {input_file}")
            
            # 验证输入文件格式
            if not self.validate_fasta(input_file):
                raise InvalidFastaFormatError(f"无效的FASTA格式: {input_file}")
            
            # 获取文件名（不含路径和扩展名）用于前缀
            if add_filename_prefix:
                filename_prefix = Path(input_file).stem
            else:
                filename_prefix = None
            
            # 读取并处理序列
            sequences = []
            for record in SeqIO.parse(input_file, "fasta"):
                # 如果需要添加文件名前缀
                if filename_prefix:
                    record.id = f"{filename_prefix}_{record.id}"
                    record.description = f"{filename_prefix}_{record.description}"
                
                sequences.append(record)
            
            # 写入输出文件
            SeqIO.write(sequences, output_file, "fasta")
            
            self.processed_count += len(sequences)
            print(f"成功处理 {len(sequences)} 条序列")
            return True
            
        except Exception as e:
            self.error_count += 1
            print(f"处理文件 {input_file} 时出错: {str(e)}")
            return False
    
    def batch_merge(self, input_files: List[str], output_dir: str,
                   add_filename_prefix: bool = False) -> Dict[str, any]:
        """
        批量合并多个FASTA文件
        
        参数:
            input_files (List[str]): 输入FASTA文件列表
            output_dir (str): 输出目录路径
            add_filename_prefix (bool): 是否添加文件名前缀到序列ID
            
        返回:
            Dict[str, any]: 处理结果统计信息
            
        异常:
            PermissionError: 输出目录无写入权限
        """
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        results = {
            'total': len(input_files),
            'success': 0,
            'failed': 0,
            'failed_files': [],
            'processed_sequences': 0
        }
        
        print(f"开始批量处理 {len(input_files)} 个文件...")
        
        for input_file in input_files:
            try:
                # 生成输出文件名
                input_path = Path(input_file)
                output_file = os.path.join(output_dir, f"merged_{input_path.name}")
                
                # 处理单个文件
                success = self.merge_sequences(input_file, output_file, add_filename_prefix)
                
                if success:
                    results['success'] += 1
                    results['processed_sequences'] += self.processed_count
                else:
                    results['failed'] += 1
                    results['failed_files'].append(input_file)
                    
            except Exception as e:
                results['failed'] += 1
                results['failed_files'].append(input_file)
                print(f"处理文件 {input_file} 时出错: {str(e)}")
        
        print("\n处理完成:")
        print(f"  总文件数: {results['total']}")
        print(f"  成功: {results['success']}")
        print(f"  失败: {results['failed']}")
        
        if results['failed_files']:
            print(f"  失败文件: {results['failed_files']}")
        
        return results
    
    def validate_fasta(self, file_path: str) -> bool:
        """
        验证FASTA文件格式是否正确
        
        参数:
            file_path (str): 要验证的FASTA文件路径
            
        返回:
            bool: 格式正确返回True，错误返回False
        """
        try:
            # 尝试解析FASTA文件
            sequences = list(SeqIO.parse(file_path, "fasta"))
            
            # 检查是否有序列
            if len(sequences) == 0:
                print(f"警告: 文件 {file_path} 中没有找到序列")
                return False
            
            # 检查序列ID是否重复
            sequence_ids = [record.id for record in sequences]
            if len(sequence_ids) != len(set(sequence_ids)):
                print(f"警告: 文件 {file_path} 中存在重复的序列ID")
                return False
            
            return True
            
        except Exception as e:
            print(f"验证文件 {file_path} 格式时出错: {str(e)}")
            return False
    
    def get_file_stats(self, file_path: str) -> Dict[str, any]:
        """
        获取FASTA文件的统计信息
        
        参数:
            file_path (str): FASTA文件路径
            
        返回:
            Dict[str, any]: 文件统计信息
        """
        try:
            sequences = list(SeqIO.parse(file_path, "fasta"))
            
            stats = {
                'total_sequences': len(sequences),
                'total_length': sum(len(record.seq) for record in sequences),
                'avg_length': 0,
                'min_length': 0,
                'max_length': 0,
                'sequence_ids': [record.id for record in sequences]
            }
            
            if sequences:
                lengths = [len(record.seq) for record in sequences]
                stats['avg_length'] = sum(lengths) / len(lengths)
                stats['min_length'] = min(lengths)
                stats['max_length'] = max(lengths)
            
            return stats
            
        except Exception as e:
            print(f"获取文件 {file_path} 统计信息时出错: {str(e)}")
            return {}


def main():
    """主函数 - 命令行接口"""
    parser = argparse.ArgumentParser(
        description='FASTA文件合并工具 - 将多行序列合并为单行格式',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 基本合并
  fasta-merge -i input.fasta -o output.fasta
  
  # 批量处理多个文件
  fasta-merge -i file1.fasta file2.fasta file3.fasta -d output_directory
  
  # 添加文件名前缀到序列ID
  fasta-merge -i *.fasta -d merged_output --add-filename-prefix
  
  # 验证FASTA文件格式
  fasta-merge -i input.fasta --validate-only
        """
    )
    
    # 输入文件参数
    parser.add_argument('-i', '--input', 
                       nargs='+', 
                       required=True,
                       help='输入FASTA文件路径（支持多个文件）')
    
    # 输出选项
    output_group = parser.add_mutually_exclusive_group(required=False)
    output_group.add_argument('-o', '--output', 
                             help='输出文件路径（单文件处理）')
    output_group.add_argument('-d', '--output-dir', 
                             help='输出目录路径（批量处理）')
    
    # 功能选项
    parser.add_argument('--add-filename-prefix', 
                       action='store_true',
                       help='添加文件名前缀到序列ID')
    parser.add_argument('--validate-only', 
                       action='store_true',
                       help='仅验证文件格式，不进行处理')
    
    # 其他选项
    parser.add_argument('--stats', 
                       action='store_true',
                       help='显示文件统计信息')
    
    args = parser.parse_args()
    
    # 创建合并器实例
    merger = FastaMerger()
    
    try:
        # 验证模式
        if args.validate_only:
            print("开始验证FASTA文件格式...")
            for file_path in args.input:
                is_valid = merger.validate_fasta(file_path)
                print(f"{file_path}: {'有效' if is_valid else '无效'}")
            return
        
        # 统计模式
        if args.stats:
            print("开始获取文件统计信息...")
            for file_path in args.input:
                stats = merger.get_file_stats(file_path)
                if stats:
                    print(f"\n文件: {file_path}")
                    print(f"  序列数量: {stats['total_sequences']}")
                    print(f"  总长度: {stats['total_length']}")
                    print(f"  平均长度: {stats['avg_length']:.1f}")
                    print(f"  最小长度: {stats['min_length']}")
                    print(f"  最大长度: {stats['max_length']}")
            return
        
        # 处理模式
        print("开始处理FASTA文件...")
        
        # 单文件处理
        if args.output and len(args.input) == 1:
            input_file = args.input[0]
            success = merger.merge_sequences(input_file, args.output, 
                                           args.add_filename_prefix)
            
            if success:
                print(f"处理完成！输出文件: {args.output}")
            else:
                print("处理失败！")
                sys.exit(1)
        
        # 批量处理
        elif args.output_dir:
            results = merger.batch_merge(args.input, args.output_dir, 
                                        args.add_filename_prefix)
            
            if results['failed'] > 0:
                print(f"有 {results['failed']} 个文件处理失败")
                sys.exit(1)
            else:
                print("所有文件处理成功！")
        
        # 参数错误
        else:
            print("错误: 必须指定 -o （单文件输出）或 -d （批量输出目录）")
            parser.print_help()
            sys.exit(1)
            
    except KeyboardInterrupt:
        print("\n操作被用户中断")
        sys.exit(1)
    except Exception as e:
        print(f"发生错误: {str(e)}")
        sys.exit(1)


if __name__ == "__main__":
    main()